import sherpa_onnx
import os
from pathlib import Path

SPEECH_MODEL_DIR = "sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20"

class Transcriber():
    def __init__(self, model_dir=None, sample_rate=16000, sd=None):
        if model_dir is None:
            print(f"模型路径空缺,请指定 [ {SPEECH_MODEL_DIR} ] 模型的绝对路径")
            exit(1)
       

        if not os.path.exists(model_dir):
            print(f"没有检测到模型, 需要将 [ {SPEECH_MODEL_DIR} ] 模型位置作为参数传入.")
            exit(1)
    

        self.model_dir = model_dir
        self.encoder = self.model_dir + "/encoder-epoch-99-avg-1.onnx"
        self.decoder = self.model_dir + "/decoder-epoch-99-avg-1.onnx"
        self.joiner  = self.model_dir + "/joiner-epoch-99-avg-1.onnx"
        self.tokens = self.model_dir + "/tokens.txt"
        self.decoding_method = "greedy_search"
        self.provider = "cpu"
        self.hotwords_file = ""
        self.hotwords_score = 1.5
        self.blank_penalty = 0.0
        self.sample_rate = sample_rate
        self.recognizer = self._create_recognizer()
        self.stream = self.recognizer.create_stream()
        

    def _assert_file_existence(self, filepath: str) -> None:
        assert Path(filepath).is_file(), (
            f"File {filepath} does not exist!\n"    
        )


    def _create_recognizer(self): 
        self._assert_file_existence(self.encoder)
        self._assert_file_existence(self.decoder)
        self._assert_file_existence(self.joiner)
        self._assert_file_existence(self.tokens)
        
        recognizer = sherpa_onnx.OnlineRecognizer.from_transducer(
            tokens=self.tokens,
            encoder=self.encoder,
            decoder=self.decoder,
            joiner=self.joiner,
            num_threads=1,
            sample_rate=self.sample_rate,
            feature_dim=80,
            enable_endpoint_detection=True,
            rule1_min_trailing_silence=2.4,
            rule2_min_trailing_silence=1.2,
            rule3_min_utterance_length=300,  # it essentially disables this rule
            decoding_method=self.decoding_method,
            provider=self.provider,
            hotwords_file=self.hotwords_file,
            hotwords_score=self.hotwords_score,
            blank_penalty=self.blank_penalty)
        print("语音识别器初始化完成,可以开始收入语音:\n")
        return recognizer

    def reset(self):
        self.recognizer.reset(self.stream)

    def transcribe(self, samples):

        self.stream.accept_waveform(self.sample_rate, samples)
        while self.recognizer.is_ready(self.stream):
            self.recognizer.decode_stream(self.stream)
        
        is_endpoint = self.recognizer.is_endpoint(self.stream)
        result = self.recognizer.get_result(self.stream)
        return result, is_endpoint
